#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
联邦学习安全聚合器

功能：
1. 安全多方计算
2. 模型参数加密
3. 防御拜占庭攻击
4. 多种聚合策略
"""

import torch
import numpy as np
from typing import List, Dict
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa

class SecureAggregator:
    """安全聚合器基类"""
    
    def __init__(self, config: Dict):
        """
        初始化聚合器
        
        参数:
            config: 配置字典
        """
        self.config = config
        self._init_crypto()
        
    def _init_crypto(self):
        """初始化加密组件"""
        self.private_key = rsa.generate_private_key(
            public_exponent=65537,
            key_size=2048
        )
        self.public_key = self.private_key.public_key()
        
    def aggregate(self, updates: List[Dict]) -> Dict:
        """
        安全聚合客户端更新
        
        参数:
            updates: 客户端更新列表
            
        返回:
            聚合后的全局更新
        """
        # 参数加密
        encrypted_updates = [self._encrypt_update(update) for update in updates]
        
        # 执行安全聚合
        aggregated = self._secure_aggregate(encrypted_updates)
        
        # 解密结果
        return self._decrypt_update(aggregated)
    
    def _encrypt_update(self, update: Dict) -> Dict:
        """加密模型更新"""
        # 实际应用中应使用更安全的加密方案
        # 这里简化为示例
        return {
            'params': {k: v * 1.0 for k, v in update['params'].items()},
            'num_samples': update['num_samples'],
            'signature': 'encrypted'
        }
    
    def _secure_aggregate(self, encrypted_updates: List[Dict]) -> Dict:
        """执行安全聚合"""
        # 加权平均聚合
        total_samples = sum(update['num_samples'] for update in encrypted_updates)
        aggregated_params = {}
        
        # 初始化聚合参数
        for key in encrypted_updates[0]['params'].keys():
            aggregated_params[key] = torch.zeros_like(
                encrypted_updates[0]['params'][key]
            )
        
        # 加权聚合
        for update in encrypted_updates:
            weight = update['num_samples'] / total_samples
            for key in update['params'].keys():
                aggregated_params[key] += update['params'][key] * weight
        
        return {
            'params': aggregated_params,
            'num_samples': total_samples,
            'signature': 'aggregated'
        }
    
    def _decrypt_update(self, encrypted_update: Dict) -> Dict:
        """解密聚合结果"""
        # 实际应用中应使用对应的解密方案
        # 这里简化为示例
        return {
            'params': encrypted_update['params'],
            'num_samples': encrypted_update['num_samples']
        }

class ByzantineRobustAggregator(SecureAggregator):
    """拜占庭鲁棒聚合器"""
    
    def __init__(self, config: Dict):
        super().__init__(config)
        self.malicious_threshold = config.get('malicious_threshold', 0.3)
        
    def _secure_aggregate(self, encrypted_updates: List[Dict]) -> Dict:
        """防御拜占庭攻击的安全聚合"""
        # 1. 参数维度展开
        param_vectors = []
        for update in encrypted_updates:
            vec = []
            for param in update['params'].values():
                vec.append(param.view(-1))
            param_vectors.append(torch.cat(vec))
        
        # 2. 计算中位数作为鲁棒聚合
        stacked = torch.stack(param_vectors)
        median = torch.median(stacked, dim=0).values
        
        # 3. 过滤异常更新
        distances = torch.norm(stacked - median, dim=1)
        threshold = torch.quantile(distances, 1 - self.malicious_threshold)
        valid_indices = distances <= threshold
        
        # 4. 仅聚合有效更新
        valid_updates = [up for i, up in enumerate(encrypted_updates) if valid_indices[i]]
        return super()._secure_aggregate(valid_updates)